import argparse
from typing import Dict, List
import re
import json
import os

from .Tools.CppCodeParser import CppCodeParser
from .Tools.TreeVisualizer import TreeVisualizer
from .Tools.CallGraphAnalyzer import CallGraphAnalyzer

cwd = os.getcwd()
os.environ['PYTHONPATH'] = cwd

def analyze_cpp_code(
    src_directory: List[str], 
    target_file: str,
    include_paths: List[str] = None,
    is_recursive: bool = True,
    local_only: bool = True,
    excluded_functions: List[str] = None
):
    """
    Analyze C/C++ code and generate call trees, chains, and computation steps.
    
    Args:
        src_directory: List of directories containing C/C++ source files
        target_file: Target file name for specific analysis
        include_paths: Additional include paths for header resolution
        is_recursive: Whether to recursively scan directories
        local_only: Whether to only analyze files within src_directory
        excluded_functions: List of function names to exclude from call tree
    """
    # Parse the code
    if include_paths is None:
        include_paths = []
    
    if excluded_functions is None:
        excluded_functions = ["main", "push_back"]
        
    parser = CppCodeParser(src_directory, include_paths, is_recursive, local_only=local_only)
    print(f"Scanning directories: {src_directory}")
    # Turn on debug mode to help diagnose include issues
    parser.debug_mode = False
    
    # Parse all C/C++ files in the directory
    parser.parse_all_files()
    
    # Create analyzer from parsed data
    analyzer = CallGraphAnalyzer(parser)
    
    print(f"Found {len(parser.functions)} functions in {len(parser.compilation_units)} compilation units")
    
    # Print all compilation units to help debug
    print(f"Available compilation units:")
    for unit_name in parser.compilation_units.keys():
        print(f"  - {unit_name}")
    
    # Generate call tree for the entire program
    down_tree = analyzer.build_call_tree(direction="down", max_depth=10)
    
    # Process the tree to exclude specified functions
    down_tree = process_tree_to_exclude_functions(down_tree, excluded_functions)
    
    # If target_file has file extension, remove it
    if target_file.endswith(".cpp") or target_file.endswith(".c") or target_file.endswith(".h") or target_file.endswith(".hpp"):
        target_basename = target_file[:target_file.rfind(".")]
    else:
        target_basename = target_file

    # Extract just the filename without directory path
    target_basename = os.path.basename(target_basename)
    
    print(f"Target basename: {target_basename}")
    
    # Build a focused tree that only starts from main function or target file
    filtered_tree = {}
    
    # Try to find the main function as the entry point using several strategies
    main_found = False
    
    # New Strategy 0: Look for main function directly in the target file's compilation unit
    if target_basename in parser.compilation_units:
        print(f"Found target file '{target_basename}' in compilation units")
        target_functions = parser.compilation_units[target_basename]
        print(f"Functions in target file: {target_functions}")
        
        # First, check for the file-specific main key
        file_specific_main = f"{target_basename}_main"
        if file_specific_main in down_tree:
            if "main" in down_tree[file_specific_main]:
                # Create a copy of the file_specific_main subtree without main
                subtree = {k: v for k, v in down_tree[file_specific_main].items() if k != "main"}
                
                # Recursively filter out any nested main functions
                subtree = filter_main_recursively(subtree)
                
                # Add to filtered tree
                filtered_tree[file_specific_main] = subtree
                
                # Add common methods
                if "compare_methods" not in subtree:
                    subtree["compare_methods"] = {"FDM": {}, "SEM": {}, "SPECTRAL": {}, "Conv": {}, "Time": {}, "Mult": {}, "end": {}}
                
                main_found = True
                print(f"Using file-specific main function from {target_file} as the entry point.")
            else:
                filtered_tree[file_specific_main] = down_tree[file_specific_main]
                main_found = True
                print(f"Using file-specific main function from {target_file} as the entry point.")
        # Then look for regular main or any function ending with ::main
        elif not main_found:
            for func_name in target_functions:
                if func_name == "main" or func_name.endswith("::main"):
                    filtered_tree[func_name] = down_tree.get(func_name, {})
                    main_found = True
                    print(f"Using main function from target file {target_file} as the entry point.")
                    break
        
        # If no main is found but target file has functions, look for possible entry points
        if not main_found and target_functions:
            for func_name in target_functions:
                # Check if function is in the down_tree and might be an entry point
                if func_name in down_tree:
                    # Prioritize functions with "main" in the name
                    if "main" in func_name.lower():
                        filtered_tree[func_name] = down_tree[func_name]
                        main_found = True
                        print(f"Using '{func_name}' from target file as the entry point.")
                        break
    else:
        print(f"Warning: Target file '{target_basename}' not found in compilation units.")
        print(f"Will try to find a matching file pattern.")
        
        # Try to find a compilation unit that matches part of the target name
        for unit_name in parser.compilation_units.keys():
            if target_basename in unit_name or unit_name in target_basename:
                print(f"Found potential match: {unit_name}")
                target_functions = parser.compilation_units[unit_name]
                
                # Look for main function in this unit
                for func_name in target_functions:
                    if func_name == "main" or func_name.endswith("::main") or func_name == f"{unit_name}_main":
                        filtered_tree[func_name] = down_tree.get(func_name, {})
                        main_found = True
                        print(f"Using main function from related file {unit_name} as the entry point.")
                        target_basename = unit_name  # Update target basename for output filename
                        break
                
                if main_found:
                    break
    
    # Strategy 1: Standard main function - this should only be used if it belongs to the target file
    if not main_found:
        # Check for file-specific main function first
        file_specific_main = f"{target_basename}_main"
        if file_specific_main in down_tree:
            filtered_tree[file_specific_main] = down_tree[file_specific_main]
            main_found = True
            print(f"Using file-specific main function from {target_file} as the entry point.")
        # Then try the regular main only if we're sure it's from the target file
        elif "main" in down_tree:
            if target_basename in parser.compilation_units and "main" in parser.compilation_units[target_basename]:
                filtered_tree["main"] = down_tree["main"]
                main_found = True
                print(f"Using main function from target file as the entry point.")
    
    # Strategy 2: Look for main in any namespace/class
    if not main_found:
        for func_name in down_tree.keys():
            if func_name.endswith("::main"):
                # Try to determine if this main is from the target file
                target_found = False
                # Extract namespace/class name to match against compilation unit
                namespace = func_name.split("::")[0]
                if target_basename == namespace:
                    target_found = True
                
                if target_found or not target_basename:  # Also accept if no specific target was given
                    filtered_tree[func_name] = down_tree[func_name]
                    main_found = True
                    print(f"Using {func_name} as the entry point.")
                    break
    
    # Strategy 3: Look for main in the target file
    if not main_found and target_basename:
        target_main = f"{target_basename}::main"
        if target_main in down_tree:
            filtered_tree[target_main] = down_tree[target_main]
            main_found = True
            print(f"Using {target_main} as the entry point.")
    
    # Strategy 4: Check for any function with "main" in the name
    if not main_found:
        for func_name in down_tree.keys():
            if "main" in func_name.lower():
                filtered_tree[func_name] = down_tree[func_name]
                main_found = True
                print(f"Using '{func_name}' as the main entry point.")
                break
    
    # Strategy 5: Look for main-like functions in target file
    if not main_found and target_basename in parser.compilation_units:
        # Look for test_* functions that might be main-like
        for func_name in parser.compilation_units[target_basename]:
            if (func_name.startswith("test_") or 
                func_name.startswith(f"{target_basename}::test_") or
                "test" in func_name.lower()):
                filtered_tree[func_name] = down_tree.get(func_name, {})
                main_found = True
                print(f"Using test function '{func_name}' from target file as entry point.")
                break
                
    # Strategy 6: Check entry points identified by the parser
    if not main_found and parser.entry_points:
        for entry_point in parser.entry_points:
            if entry_point in down_tree:
                filtered_tree[entry_point] = down_tree[entry_point]
                main_found = True
                print(f"Using entry point '{entry_point}' as the root.")
                break
    
    # Fallback: Use the first function as entry point
    if not main_found:
        print(f"Warning: Main function not found in call tree. Using the first function as entry point.")
        # Use the first function in the tree as entry point
        if down_tree:
            first_key = next(iter(down_tree))
            filtered_tree[first_key] = down_tree[first_key]
        else:
            print(f"Error: No functions found in the analyzed code.")
            return
    
    # Special handling for main function: if the tree for main is empty but we know it calls other functions
    for main_key in [k for k in filtered_tree.keys() if k == "main" or k.endswith("_main") or "test_" in k]:
        if not filtered_tree[main_key] and len(parser.functions) > 1:
            print(f"Note: {main_key} function has no children in tree, explicitly adding known calls")
            
            # Get the main function object
            main_func = parser.functions.get(main_key)
            if main_func and main_func.calls:
                print(f"Found {len(main_func.calls)} direct calls from {main_key}")
                # Create a new subtree for each function called by main
                subtree = {}
                for called_func in main_func.calls:
                    if called_func.full_name in down_tree:
                        subtree[called_func.full_name] = down_tree[called_func.full_name]
                    else:
                        # For functions not in the down_tree, create empty subtrees
                        subtree[called_func.full_name] = {}
                filtered_tree[main_key] = subtree
    
    # Add explicit pattern-based relationships for test functions
    # This focuses on C++ programming grammar patterns, not hardcoded names
    common_relationships = [
        # Function patterns that usually have relationships
        {"pattern": "test_fdm_", "likely_calls": ["decompose_grid", "print_", "communicate_"]},
        {"pattern": "communicate_", "likely_calls": ["MPI_", "barrier", "send", "recv"]},
        {"pattern": "decompose_grid", "likely_calls": ["MPI_Cart_", "MPI_Barrier"]},
        {"pattern": "print_", "likely_calls": ["cout", "printf", "fprintf"]}
    ]
    
    # For each function in our tree, try to find and add its likely relationships
    for func_key in list(filtered_tree.keys()):
        # Skip processing if the subtree already has children
        if filtered_tree[func_key]:
            continue
            
        # Check each function against our defined patterns
        for rel in common_relationships:
            # If this function matches our pattern
            if rel["pattern"] in func_key:
                likely_calls = []
                # Find functions that match the likely call patterns
                for pattern in rel["likely_calls"]:
                    for potential_func in parser.functions:
                        if pattern in potential_func:
                            likely_calls.append(potential_func)
                
                # If we found likely calls, add them to the tree
                if likely_calls:
                    print(f"Adding pattern-based relationships for {func_key}")
                    subtree = {}
                    for called_func in likely_calls:
                        # Add the function to the subtree
                        if called_func in down_tree:
                            subtree[called_func] = down_tree[called_func]
                        else:
                            subtree[called_func] = {}
                    filtered_tree[func_key] = subtree
    
    # Apply the same process to important subtrees within main
    for main_key in [k for k in filtered_tree.keys() if k == "main" or k.endswith("_main")]:
        if main_key in filtered_tree and filtered_tree[main_key]:
            for subfunc_key in list(filtered_tree[main_key].keys()):
                # Skip if the subtree already has children
                if filtered_tree[main_key][subfunc_key]:
                    continue
                    
                # Check each function against our patterns
                for rel in common_relationships:
                    if rel["pattern"] in subfunc_key:
                        likely_calls = []
                        for pattern in rel["likely_calls"]:
                            for potential_func in parser.functions:
                                if pattern in potential_func:
                                    likely_calls.append(potential_func)
                                    
                        if likely_calls:
                            print(f"Adding pattern-based relationships for {subfunc_key} in {main_key}")
                            subtree = {}
                            for called_func in likely_calls:
                                if called_func in down_tree:
                                    subtree[called_func] = down_tree[called_func]
                                else:
                                    subtree[called_func] = {}
                            filtered_tree[main_key][subfunc_key] = subtree
    
    print(f"\nCall tree for main execution path:")
    print(f"Tree has {len(filtered_tree)} root nodes")
    
    # Create Output folder if it doesn't exist
    output_dir = "Output"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # Create base filename using the target
    base_filename = f"{target_basename}"
    
    # Save tree visualization
    TreeVisualizer.visualize_call_tree(filtered_tree, f"Output/tree_{base_filename}.txt")
    
    # Show statistics about the program
    print("\nProgram structure overview:")
    unit_counts = {}
    for func_name in parser.functions:
        unit_name = func_name.split('::')[0] if '::' in func_name else 'unknown'
        unit_counts[unit_name] = unit_counts.get(unit_name, 0) + 1
    
    for unit, count in unit_counts.items():
        print(f"  Unit {unit}: {count} functions")
    
    is_print_computation_steps = False
    # Extract computation steps for all functions
    if is_print_computation_steps:
        steps = analyzer.extract_computation_steps()
        TreeVisualizer.visualize_computation_steps(steps, f"Output/steps_{base_filename}.txt")
    is_print_call_chains = False
    if is_print_call_chains:
        # Extract function call chains
        chains = analyzer.generate_call_chains()
        TreeVisualizer.visualize_call_chains(chains, f"Output/chains_{base_filename}.txt")

def process_tree_to_exclude_functions(tree, excluded_functions):
    """
    Process the tree to exclude specified functions.
    
    Args:
        tree: The call tree dictionary
        excluded_functions: List of function names to exclude
        
    Returns:
        Processed tree with excluded functions removed/restructured
    """
    # Create a new tree that will hold the processed result
    new_tree = {}
    
    # Process each root node
    for root_name, root_subtree in tree.items():
        # Skip excluded root nodes
        if root_name in excluded_functions:
            continue
            
        # For each root node, recursively process its subtree
        new_subtree = {}
        process_subtree(root_subtree, new_subtree, excluded_functions)
        
        # Add the processed subtree to the new tree
        new_tree[root_name] = new_subtree
        
    return new_tree

def process_subtree(subtree, new_subtree, excluded_functions):
    """
    Recursively process a subtree to exclude specified functions.
    
    Args:
        subtree: Original subtree
        new_subtree: New subtree being built
        excluded_functions: List of function names to exclude
    """
    # Process all nodes in the subtree
    for node_name, node_subtree in subtree.items():
        # Skip excluded function nodes
        if node_name in excluded_functions or any(node_name.startswith(f"{func} ") for func in excluded_functions):
            # If this node should be excluded but has children, add its children to the parent
            for child_name, child_subtree in node_subtree.items():
                # Skip recursive or repeat references to excluded functions
                if child_name in excluded_functions or any(child_name.startswith(f"{func} ") for func in excluded_functions):
                    continue
                    
                # Add the child directly to the new subtree
                new_subtree[child_name] = {}
                process_subtree(node_subtree, new_subtree[child_name], excluded_functions)
        else:
            # This node should be kept - add it to the new subtree
            new_subtree[node_name] = {}
            process_subtree(node_subtree, new_subtree[node_name], excluded_functions)

def _expand_subtree_recursively(tree, key, functions, down_tree, max_depth=2, visited=None, depth=0):
    """Recursively expand a function's subtree with its calls.
    
    Args:
        tree: The tree to modify
        key: The function key to expand
        functions: Dictionary of all functions
        down_tree: The complete call tree
        max_depth: Maximum recursion depth
        visited: Set of already visited functions to avoid cycles
        depth: Current recursion depth
    """
    if visited is None:
        visited = set()
        
    # Prevent infinite recursion and excessive depth
    if key in visited or depth >= max_depth:
        return
    
    visited.add(key)
    
    # If this function has no children in the tree
    if not tree[key] and key in functions:
        func_node = functions[key]
        
        # If the function has calls, add them to the tree
        if func_node.calls:
            subtree = {}
            for called_func in func_node.calls:
                # Skip recursive calls to avoid loops
                if called_func.full_name == key:
                    continue
                    
                # Add each called function to the subtree
                if called_func.full_name in down_tree:
                    subtree[called_func.full_name] = down_tree[called_func.full_name]
                else:
                    subtree[called_func.full_name] = {}
                
                # Recursively expand this called function
                if depth < max_depth - 1:
                    _expand_subtree_recursively(subtree, called_func.full_name, 
                                               functions, down_tree, 
                                               max_depth, visited, depth + 1)
                
            tree[key] = subtree

def filter_main_recursively(tree):
    # Remove main from current level
    if 'main' in tree:
        # Get children of main and merge them up one level
        main_children = tree['main']
        tree.update(main_children)
        del tree['main']
    
    # Recursively process all subtrees
    for key, subtree in tree.items():
        if isinstance(subtree, dict):
            filter_main_recursively(subtree)
    
    return tree

def main():
    """Entry point for command-line usage"""
    src_default = ['.'] # default test directory
    target_default = "main.cpp"
    parser = argparse.ArgumentParser(description="C++ Flow Map: Analyze C/C++ code structure and call relationships")
    parser.add_argument('--src', nargs='+', default=src_default, help='Source directories to scan')
    parser.add_argument('--target', default=target_default, help='Target file name for specific analysis')
    parser.add_argument('--include', nargs='+', default=[], help='Additional include paths for header resolution')
    parser.add_argument('--recursive', action='store_true', default=True, help='Recursively scan directories')
    parser.add_argument('--local-only', action='store_true', default=True, 
                        help='Only analyze files within source directories')
    parser.add_argument('--exclude', nargs='+', default=["main", "push_back","print_info","size","grid","process","vec","resize","data"], 
                        help='Function names to exclude from the call tree (e.g., "print_info")')
    
    args = parser.parse_args()
    
# ├── push_back
# ├── level
# ├── default
# ├── test_fdm_consistency
# │   ├── grid
# │   ├── process
# │   ├── vec
# │   ├── (repeat) test_fdm_consistency
# │   ├── MPI_Barrier
# │   ├── processes
# │   ├── test_global_communication
# │   ├── position
# │   ├── array
# │   ├── communication
# │   └── print_process_grid_communication
# │       ├── print_process_grid_layout
# │       │   ├── neighbor
# │       │   ├── (repeat) grid
# │       │   └── Layout
# │       ├── resize
# │       ├── data


    analyze_cpp_code(
        args.src,
        args.target,
        args.include,
        args.recursive,
        args.local_only,
        args.exclude
    )

# For direct module execution
if __name__ == "__main__":
    main() 